"""
python fujian2_RandomForest_predict_tqdm_handle5less.py
"""
import os
import re
import json
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from tqdm import tqdm


def read_json_data(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df['sales'] = df['sales'].astype(int)
    return df

def prepare_data(df):
    # 对日期进行排序
    df = df.sort_values('date')
    # 提取特征和目标
    X = df.index.values.reshape(-1, 1)
    y = df['sales'].values
    return X, y

def train_random_forest(X_train, y_train):
    rf = RandomForestRegressor(n_estimators=100, random_state=42)
    rf.fit(X_train, y_train)
    return rf

def predict_future_sales(model, last_date, num_days):
    future_dates = [last_date + timedelta(days=i) for i in range(1, num_days + 1)]
    future_dates_df = pd.DataFrame({'date': future_dates})
    
    # 确保索引的计算正确
    future_dates_df['index'] = future_dates_df.index.values.reshape(-1, 1)
    X_future = future_dates_df['index'].values.reshape(-1, 1)
    
    predictions = model.predict(X_future)
    return future_dates_df, predictions

def fill_with_average(df, num_days):
    # 计算均值
    avg_sales = df['sales'].mean()
    # 生成未来的日期
    last_date = df['date'].max()
    future_dates = [last_date + timedelta(days=i) for i in range(1, num_days + 1)]
    # 生成预测结果
    return pd.DataFrame({'date': future_dates}), [avg_sales] * num_days

def save_predictions_to_json(category_id, predictions, output_folder):
    future_dates, predicted_sales = predictions
    result = [{'date': date.strftime('%Y/%m/%d'), 'predicted_sales': int(sale)} for date, sale in zip(future_dates['date'], predicted_sales)]
    output_json_path = os.path.join(output_folder, f'json_category{category_id}.json')
    os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
    with open(output_json_path, 'w') as f:
        json.dump(result, f)

def plot_data_and_predictions(df, predictions, category_id, output_folder):
    original_dates = df['date']
    original_sales = df['sales']
    future_dates, predicted_sales = predictions
    plt.figure(figsize=(10, 6))
    plt.scatter(original_dates, original_sales, label='Original Data', c='blue')
    plt.scatter(future_dates['date'], predicted_sales, label='Predicted Data', c='red')
    
    plt.legend()
    plt.xlabel('Date')
    plt.ylabel('Sales')
    plt.title(f'Random Forest Prediction on Category {category_id}')
    output_image_path = os.path.join(output_folder, 'gragh', f'gragh_category{category_id}.png')
    os.makedirs(os.path.dirname(output_image_path), exist_ok=True)
    plt.savefig(output_image_path)
    plt.close()  # 关闭当前图形

if __name__ == "__main__":
    input_folder = r'..\fujian\fujian2\groupByCategory'
    output_folder = r'..\fujian\fujian2\RandomForest'
    
    json_files = [os.path.join(input_folder, file) for file in os.listdir(input_folder) if file.endswith('.json')]
    
    for json_file in tqdm(json_files, desc="Processing categories"):
        filename = os.path.basename(json_file)
        
        # 提取 category ID
        category_id = -1
        match = re.search(r'category_(?:category)?(\d+)', filename)
        if match:
            category_id = int(match.group(1))
        else:
            print("未找到匹配的 category ID")
            continue
        
        # 读取数据
        df = read_json_data(json_file)
        
        # 检查样本数量是否足够
        if len(df) < 5:
            print(f"数据样本不足: {len(df)}, category的id是: ", category_id)
            # 样本不足时，使用均值进行填充
            future_predictions = fill_with_average(df, 92)
        else:
            # 准备数据
            X, y = prepare_data(df)
            # 划分训练集和测试集
            split_index = int(0.8 * len(df))
            X_train, y_train = X[:split_index], y[:split_index]
            
            # 训练随机森林模型
            model = train_random_forest(X_train, y_train)
            
            # 预测未来销售
            last_date = df['date'].max()
            future_predictions = predict_future_sales(model, last_date, 92)
        
        # 保存预测到 JSON 文件
        save_predictions_to_json(category_id, future_predictions, output_folder)
        
        # 绘制图形
        plot_data_and_predictions(df, future_predictions, category_id, output_folder)
